#!/usr/bin/env python
# -*- coding: utf-8 -*-


import os
import glob, sys
import numpy as np
import xarray as xr
import pandas as pd


def fmc_calculation(lfm_h,lfm_w,GSI):
    min_w = 60.
    min_h = 30.
    max_w = 200.
    max_h = 250.
    GU_h = 0.5
    GU_w = 0.5

    m_h = (max_h-min_h)/(1.-GU_h)
    b_h = max_h-m_h
    m_w = (max_w-min_w)/(1.-GU_w)
    b_w = max_w-m_w

    lfm_h = xr.where(GSI>=GU_h,m_h*GSI+b_h,lfm_h)
    lfm_w = xr.where(GSI>=GU_w,m_w*GSI+b_w,lfm_w)
    lfm_h = xr.where(lfm_h>max_h,max_h,lfm_h)
    lfm_w = xr.where(lfm_w>max_w,max_w,lfm_w)

    return lfm_h, lfm_w

year=int(sys.argv[1])     # Enter year
month=int(sys.argv[2])    # Enter month
month_str = ['00', '01', '02', '03', '04', '05', '06', '07', '08', '09', '10', '11', '12']

model=str(sys.argv[3])
experiment=str(sys.argv[4])

path_gsi='../OUTS/GSI/'+model+'/'+experiment+'/means/'
ds_gsi = xr.open_dataset(path_gsi+'ERA5_gsi_{}-{}.nc'.format(year,month_str[month]))     

path_out_fmc='../OUTS/LFMC/'+model+'/'+experiment+'/'

os.system('mkdir -p '+path_out_fmc+'Herb')
os.system('mkdir -p '+path_out_fmc+'Woody')

ds_fmch = xr.full_like(ds_gsi,fill_value=30.)
ds_fmcw = xr.full_like(ds_gsi,fill_value=60.)

lfm_h, lfm_w = fmc_calculation(ds_fmch.GSI,ds_fmcw.GSI,ds_gsi.GSI)

ds_fmch.GSI.values[:] = lfm_h.copy()
ds_fmch = ds_fmch['GSI'].rename('fmc')
ds_fmch.attrs['standard_name'] = 'fmc'
ds_fmch.attrs['long_name'] = 'Herb live fuel moisture content'
ds_fmch.attrs['units'] = '%'
ds_fmch.to_netcdf(path_out_fmc+'Herb/ERA5_fmch_{}-{}.nc'.format(year,month_str[month]))

ds_fmcw.GSI.values[:] = lfm_w.values.copy()
ds_fmcw = ds_fmcw['GSI'].rename('fmc')
ds_fmcw.attrs['standard_name'] = 'fmc'
ds_fmcw.attrs['long_name'] = 'Woody live fuel moisture content'
ds_fmcw.attrs['units'] = '%'
ds_fmcw.to_netcdf(path_out_fmc+'Woody/ERA5_fmcw_{}-{}.nc'.format(year,month_str[month]))
